#include "tensorRTTools.h"
#include <opencv2/opencv.hpp>
#include <queue>
#include <thread>
#include <mutex>
#include <condition_variable>

#include "cpm.hpp"
#include "infer.hpp"
#include "yolo.hpp"


void to_upper(std::string& str)
{
	for (char& c : str) {
		c = std::toupper(static_cast<unsigned char>(c));
	}
}

tensorRTTools::tensorRTTools(std::string enginePath, std::map<int, std::string>maplabels, std::string strversion)
	: m_maplabels(maplabels), m_stopThread(false) {  // 初始化标签列表和线程控制标志


	to_upper(strversion);
	   // 根据 strversion 设置 YOLO 类型
	yolo::Type yoloType;
	if (strversion == "V5") {
		yoloType = yolo::Type::V5;
	}
	else if (strversion == "X") {
		yoloType = yolo::Type::X;
	}
	else if (strversion == "V3") {
		yoloType = yolo::Type::V3;
	}
	else if (strversion == "V7") {
		yoloType = yolo::Type::V7;
	}
	else if (strversion == "V8") {
		yoloType = yolo::Type::V8;
	}
	else if (strversion == "V8SEG") {
		yoloType = yolo::Type::V8Seg;
	}
	else {
		throw std::invalid_argument("Unsupported YOLO version: " + strversion);
	}
	m_yolo = yolo::load(enginePath, yoloType);
	if (m_yolo == nullptr) {
		throw std::runtime_error("Failed to load YOLO engine.");
	}

	// 启动检测线程
	m_detectionThread = std::thread(&tensorRTTools::detectionThreadFunc, this);
}

tensorRTTools::~tensorRTTools() {
	{
		// 设置停止标志并通知检测线程
		std::lock_guard<std::mutex> lock(m_mutex);
		while (!m_imageQueue.empty())
			m_imageQueue.pop(); 
		m_stopThread = true;
		m_condition.notify_all();
	}

	// 等待线程结束
	if (m_detectionThread.joinable()) {
		m_detectionThread.join();
	}
}

void tensorRTTools::detector(const cv::Mat& matRGB24) {
	if (matRGB24.empty()) {
		throw std::invalid_argument("Input image is empty.");
	}

	// 将图像推入队列
	{
		std::lock_guard<std::mutex> lock(m_mutex);
		if(m_imageQueue.size()>3)
		{
			while (!m_imageQueue.empty())
				m_imageQueue.pop();
		}
		m_imageQueue.push(matRGB24.clone());  // 使用clone防止图像数据被外部修改
	}

	// 通知检测线程
	m_condition.notify_one();
}

void tensorRTTools::detectionThreadFunc() {
	while (!m_stopThread.load()) {
		cv::Mat image;

		{
			std::unique_lock<std::mutex> lock(m_mutex);

			// 等待队列有数据或线程停止
			m_condition.wait(lock, [this] { return !m_imageQueue.empty() || m_stopThread; });

			if (m_stopThread && m_imageQueue.empty()) {
				break;  // 退出线程
			}

			// 从队列中取出一张图像
			image = m_imageQueue.front();
			m_imageQueue.pop();
		}

		// 执行YOLO检测
		auto start_time = std::chrono::high_resolution_clock::now();

		// 将cv::Mat 转换为 yolo::Image
		yolo::Image yoloImage = yolo::Image(image.data, image.cols, image.rows);
		yolo::BoxArray detections;
		detections = m_yolo->forward(yoloImage);	
		{
			std::unique_lock<std::mutex> lock(m_mutex);
			m_detections = detections;		
		}
		// 如果回调函数已设置，则调用它
		if (m_detectionCallback) {
			m_detectionCallback(m_detections);
		}

		std::this_thread::sleep_for(std::chrono::milliseconds(50));

	}
}

void tensorRTTools::drawRectangle(cv::Mat& matRGB24) {
	yolo::BoxArray detections;
	{
		std::unique_lock<std::mutex> lock(m_mutex);
		detections = m_detections;
	}

	for (auto& obj : detections) {
		uint8_t b, g, r;
		std::tie(b, g, r) = yolo::random_color(obj.class_label);

		int rectangle_thickness = 2;

		cv::rectangle(matRGB24, cv::Point(obj.left, obj.top), cv::Point(obj.right, obj.bottom),
			cv::Scalar(b, g, r), rectangle_thickness);

		auto name = m_maplabels[obj.class_label];
		auto caption = cv::format("%s %.2f", name.c_str(), obj.confidence);
	
		double font_scale = 0.5;
		int font_thickness = 1;

		int width = cv::getTextSize(caption, cv::FONT_HERSHEY_SIMPLEX, font_scale, font_thickness, nullptr).width + 10;
		cv::rectangle(matRGB24, cv::Point(obj.left - 3, obj.top - 22),
			cv::Point(obj.left + width, obj.top), cv::Scalar(b, g, r), -1);

		cv::putText(matRGB24, caption, cv::Point(obj.left, obj.top - 5), cv::FONT_HERSHEY_SIMPLEX,
			font_scale, cv::Scalar::all(0), font_thickness, cv::LINE_AA);
	}
}

void tensorRTTools::setDetectionCallback(DetectionCallback callback) {
	m_detectionCallback = std::move(callback);
}
